package com.whz.gateway.limit;

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.support.ipresolver.RemoteAddressResolver;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author whz
 */
@Slf4j
public class GatewayRemoteIpResolver implements RemoteAddressResolver {

    public static final String BEAN_NAME = "gatewayRemoteIpResolver";

    private static final String[] HEADERS = new String[] {
            "x-forwarded-for",
            "Proxy-Client-IP",
            "WL-Proxy-Client-IP",
            "HTTP_CLIENT_IP",
            "X-Real-IP",
    };

    public static final String UNKNOWN = "unknown";

    @Override
    public InetSocketAddress resolve(ServerWebExchange exchange) {
        InetSocketAddress ip = getRemoteIpAddrByRequest(exchange.getRequest());
        log.info("Get remote ip:{}", ip);
        return ip;
    }

    public InetSocketAddress getRemoteIpAddrByRequest(final ServerHttpRequest request) {
        List<String> ip;
        for (String header : HEADERS) {
            ip = checkNull(request.getHeaders().get(header));
            if (CollectionUtils.isEmpty(ip)) {
                continue;
            }
            log.info("Get remote ip:{} from request-header:{}", ip, header);
            return new InetSocketAddress(StringUtils.trimWhitespace(ip.get(0)), 0);
        }
        return request.getRemoteAddress();
    }

    private static List<String> checkNull(List<String> ips) {
        return CollectionUtils.isEmpty(ips)? Lists.newArrayList() :
                ips.stream().filter(ip -> !check(ip)).collect(Collectors.toList());
    }

    private static boolean check(String ip) {
        return StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(StringUtils.trimWhitespace(ip));
    }
}